//! Generic monomorphization utilities.
//!
//! This module introduces the infrastructure required to turn generic
//! declarations into concrete, monomorphic declarations. It does not yet wire
//! into the full compilation pipeline, but it provides the caching, template
//! bookkeeping, and type substitution helpers needed by the upcoming passes.

use crate::ast::decl::{Declaration, FunctionDecl, GenericParam, StructDecl};
use crate::ast::types::*;
use crate::ast::Program;
use crate::error::{CompilerError, Result};
use crate::utils::mangling;
use std::collections::HashMap;

/// Key that identifies a specific instantiation of a template.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct InstantiationKey {
    template_name: String,
    type_args: Vec<Type>,
}

impl InstantiationKey {
    fn new(name: impl Into<String>, type_args: &[Type]) -> Self {
        Self {
            template_name: name.into(),
            type_args: type_args.to_vec(),
        }
    }
}

/// Stores all generic templates discovered in the AST and handles requests to
/// instantiate them with specific type arguments.
#[derive(Default)]
pub struct Monomorphizer {
    function_templates: HashMap<String, Vec<FunctionDecl>>, // keep overloads ordered
    struct_templates: HashMap<String, Vec<StructDecl>>,
    instantiated_functions: HashMap<InstantiationKey, FunctionDecl>,
    instantiated_structs: HashMap<InstantiationKey, StructDecl>,
    pending_function_keys: Vec<InstantiationKey>,
    pending_struct_keys: Vec<InstantiationKey>,
}

impl Monomorphizer {
    /// Create a new monomorphizer.
    pub fn new() -> Self {
        Self::default()
    }

    /// Scan the program and cache all generic templates. This should be called
    /// before attempting any instantiation so that templates are available.
    pub fn collect_templates(&mut self, program: &Program) {
        for declaration in &program.declarations {
            match declaration {
                Declaration::Function(func) if !func.generics.is_empty() => {
                    self.function_templates
                        .entry(func.name.clone())
                        .or_default()
                        .push(func.clone());
                }
                Declaration::Struct(struct_decl) if !struct_decl.generics.is_empty() => {
                    self.struct_templates
                        .entry(struct_decl.name.clone())
                        .or_default()
                        .push(struct_decl.clone());
                }
                _ => {}
            }
        }
    }

    /// Instantiate the newest function template registered under `name` with
    /// the provided type arguments.
    pub fn instantiate_function(&mut self, name: &str, type_args: &[Type]) -> Result<FunctionDecl> {
        let template = self
            .function_templates
            .get(name)
            .ok_or_else(|| {
                CompilerError::semantic(0, 0, format!("Generic function '{}' not found", name))
            })?
            .last()
            .cloned()
            .expect("templates Vec should be non-empty");

        self.instantiate_function_from_template(&template, type_args)
    }

    /// Instantiate the newest struct template registered under `name` with the
    /// provided type arguments.
    pub fn instantiate_struct(&mut self, name: &str, type_args: &[Type]) -> Result<StructDecl> {
        let template = self
            .struct_templates
            .get(name)
            .ok_or_else(|| {
                CompilerError::semantic(0, 0, format!("Generic struct '{}' not found", name))
            })?
            .last()
            .cloned()
            .expect("templates Vec should be non-empty");

        self.instantiate_struct_from_template(&template, type_args)
    }

    fn instantiate_function_from_template(
        &mut self,
        template: &FunctionDecl,
        type_args: &[Type],
    ) -> Result<FunctionDecl> {
        if template.generics.len() != type_args.len() {
            return Err(CompilerError::semantic(
                template.location.line,
                template.location.column,
                format!(
                    "Function '{}' expects {} type arguments but {} were provided",
                    template.name,
                    template.generics.len(),
                    type_args.len()
                ),
            ));
        }

        let key = InstantiationKey::new(&template.name, type_args);
        if let Some(existing) = self.instantiated_functions.get(&key) {
            return Ok(existing.clone());
        }

        let substitution = build_substitution_map(&template.generics, type_args);
        let mut instantiated = template.clone();
        instantiated.generics.clear();
        instantiated.name = Self::mangle_instance_name(&template.name, type_args);

        for param in &mut instantiated.parameters {
            param.param_type = substitute_type(&param.param_type, &substitution);
        }

        if let Some(return_type) = &mut instantiated.return_type {
            *return_type = substitute_type(return_type, &substitution);
        }

        // NOTE: We intentionally keep the function body untouched for now. A
        // future patch will walk statements/expressions and rewrite embedded
        // type annotations. Keeping the body avoids losing user code while the
        // remaining passes are wired up.

        self.instantiated_functions
            .insert(key.clone(), instantiated.clone());
        self.pending_function_keys.push(key);

        Ok(instantiated)
    }

    fn instantiate_struct_from_template(
        &mut self,
        template: &StructDecl,
        type_args: &[Type],
    ) -> Result<StructDecl> {
        if template.generics.len() != type_args.len() {
            return Err(CompilerError::semantic(
                template.location.line,
                template.location.column,
                format!(
                    "Struct '{}' expects {} type arguments but {} were provided",
                    template.name,
                    template.generics.len(),
                    type_args.len()
                ),
            ));
        }

        let key = InstantiationKey::new(&template.name, type_args);
        if let Some(existing) = self.instantiated_structs.get(&key) {
            return Ok(existing.clone());
        }

        let substitution = build_substitution_map(&template.generics, type_args);
        let mut instantiated = template.clone();
        instantiated.generics.clear();
        instantiated.name = Self::mangle_instance_name(&template.name, type_args);

        for field in &mut instantiated.fields {
            field.field_type = substitute_type(&field.field_type, &substitution);
        }

        for method in &mut instantiated.methods {
            for param in &mut method.parameters {
                param.param_type = substitute_type(&param.param_type, &substitution);
            }
            if let Some(return_type) = &mut method.return_type {
                *return_type = substitute_type(return_type, &substitution);
            }
        }

        self.instantiated_structs
            .insert(key.clone(), instantiated.clone());
        self.pending_struct_keys.push(key);

        Ok(instantiated)
    }

    /// Drain the list of newly instantiated functions since the previous call.
    pub fn take_pending_function_instantiations(&mut self) -> Vec<FunctionDecl> {
        let mut decls = Vec::new();
        for key in self.pending_function_keys.drain(..) {
            if let Some(func) = self.instantiated_functions.get(&key) {
                decls.push(func.clone());
            }
        }
        decls
    }

    /// Check whether a struct template with the given name exists.
    pub fn has_struct_template(&self, name: &str) -> bool {
        self.struct_templates.contains_key(name)
    }

    /// Drain newly instantiated structs.
    pub fn take_pending_struct_instantiations(&mut self) -> Vec<StructDecl> {
        let mut decls = Vec::new();
        for key in self.pending_struct_keys.drain(..) {
            if let Some(struct_decl) = self.instantiated_structs.get(&key) {
                decls.push(struct_decl.clone());
            }
        }
        decls
    }

    fn mangle_instance_name(base: &str, type_args: &[Type]) -> String {
        if type_args.is_empty() {
            return format!("{}__mono", base);
        }
        let encoded: Vec<String> = type_args.iter().map(|ty| mangling::type_tag(ty)).collect();
        format!("{}__{}", base, encoded.join("_"))
    }
}

/// Build a mapping from generic parameter name to the specific type argument.
fn build_substitution_map(generics: &[GenericParam], type_args: &[Type]) -> HashMap<String, Type> {
    generics
        .iter()
        .zip(type_args.iter())
        .map(|(param, ty)| (param.name.clone(), ty.clone()))
        .collect()
}

/// Substitute occurrences of `Type::Generic` using the provided mapping.
fn substitute_type(ty: &Type, substitutions: &HashMap<String, Type>) -> Type {
    match ty {
        Type::Generic(name) => substitutions
            .get(name)
            .cloned()
            .unwrap_or_else(|| ty.clone()),
        Type::Array(array) => Type::Array(ArrayType {
            element_type: Box::new(substitute_type(&array.element_type, substitutions)),
            ..array.clone()
        }),
        Type::Map(map) => Type::Map(MapType {
            key_type: Box::new(substitute_type(&map.key_type, substitutions)),
            value_type: Box::new(substitute_type(&map.value_type, substitutions)),
            ..map.clone()
        }),
        Type::Tuple(tuple) => Type::Tuple(TupleType {
            element_types: tuple
                .element_types
                .iter()
                .map(|t| substitute_type(t, substitutions))
                .collect(),
            ..tuple.clone()
        }),
        Type::Struct(struct_type) => {
            if struct_type.type_args.is_empty() {
                if let Some(replacement) = substitutions.get(&struct_type.name) {
                    return replacement.clone();
                }
            }
            Type::Struct(StructType {
                type_args: struct_type
                    .type_args
                    .iter()
                    .map(|t| substitute_type(t, substitutions))
                    .collect(),
                ..struct_type.clone()
            })
        }
        Type::Enum(enum_type) => Type::Enum(EnumType {
            type_args: enum_type
                .type_args
                .iter()
                .map(|t| substitute_type(t, substitutions))
                .collect(),
            ..enum_type.clone()
        }),
        Type::Function(function_type) => Type::Function(FunctionType {
            parameter_types: function_type
                .parameter_types
                .iter()
                .map(|t| substitute_type(t, substitutions))
                .collect(),
            return_type: function_type
                .return_type
                .as_ref()
                .map(|t| Box::new(substitute_type(t, substitutions))),
            ..function_type.clone()
        }),
        Type::Reference(reference_type) => Type::Reference(ReferenceType {
            referenced_type: Box::new(substitute_type(
                &reference_type.referenced_type,
                substitutions,
            )),
            ..reference_type.clone()
        }),
        Type::Vec(vec_type) => Type::Vec(VecType {
            element_type: Box::new(substitute_type(&vec_type.element_type, substitutions)),
            ..vec_type.clone()
        }),
        Type::Channel(channel_type) => Type::Channel(ChannelType {
            element_type: Box::new(substitute_type(&channel_type.element_type, substitutions)),
            ..channel_type.clone()
        }),
        Type::Optional(optional_type) => Type::Optional(OptionalType {
            inner_type: Box::new(substitute_type(&optional_type.inner_type, substitutions)),
            ..optional_type.clone()
        }),
        Type::Error(error_type) => Type::Error(ErrorType {
            inner_type: Box::new(substitute_type(&error_type.inner_type, substitutions)),
            ..error_type.clone()
        }),
        Type::Result(result_type) => Type::Result(ResultType {
            ok_type: Box::new(substitute_type(&result_type.ok_type, substitutions)),
            err_type: Box::new(substitute_type(&result_type.err_type, substitutions)),
            ..result_type.clone()
        }),
        Type::Option(option_type) => Type::Option(OptionType {
            some_type: Box::new(substitute_type(&option_type.some_type, substitutions)),
            ..option_type.clone()
        }),
        Type::Union(union_type) => Type::Union(UnionType {
            member_types: union_type
                .member_types
                .iter()
                .map(|t| substitute_type(t, substitutions))
                .collect(),
            ..union_type.clone()
        }),
        Type::Alias(alias_type) => Type::Alias(AliasType {
            type_args: alias_type
                .type_args
                .iter()
                .map(|t| substitute_type(t, substitutions))
                .collect(),
            ..alias_type.clone()
        }),
        Type::Rc(rc_type) => Type::Rc(RcType {
            inner_type: Box::new(substitute_type(&rc_type.inner_type, substitutions)),
            ..rc_type.clone()
        }),
        Type::Weak(weak_type) => Type::Weak(WeakType {
            inner_type: Box::new(substitute_type(&weak_type.inner_type, substitutions)),
            ..weak_type.clone()
        }),
        Type::Basic(_) => ty.clone(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ast::decl::FunctionDecl;
    use crate::ast::types::{BasicType, Parameter};
    use crate::ast::Block;
    use crate::error::Location;

    fn generic_function_decl() -> FunctionDecl {
        FunctionDecl {
            name: "identity".to_string(),
            generics: vec![GenericParam {
                name: "T".to_string(),
                constraints: vec![],
                location: Location::new(1, 1, 0),
            }],
            parameters: vec![Parameter {
                name: "value".to_string(),
                param_type: Type::Generic("T".to_string()),
                default_value: None,
                location: Location::new(1, 1, 0),
            }],
            return_type: Some(Type::Generic("T".to_string())),
            body: Some(Block {
                statements: vec![],
                location: Location::new(1, 1, 0),
            }),
            attributes: vec![],
            location: Location::new(1, 1, 0),
        }
    }

    #[test]
    fn substitutes_function_signature() {
        let mut monomorphizer = Monomorphizer::new();
        monomorphizer
            .function_templates
            .entry("identity".to_string())
            .or_insert_with(Vec::new)
            .push(generic_function_decl());

        let instantiated = monomorphizer
            .instantiate_function("identity", &[Type::Basic(BasicType::Int)])
            .expect("instantiation should succeed");

        assert!(instantiated.generics.is_empty());
        assert_eq!(instantiated.name, "identity__int");
        assert!(matches!(
            instantiated.parameters[0].param_type,
            Type::Basic(BasicType::Int)
        ));
        assert!(matches!(
            instantiated.return_type,
            Some(Type::Basic(BasicType::Int))
        ));

        let pending = monomorphizer.take_pending_function_instantiations();
        assert_eq!(pending.len(), 1);
        assert_eq!(pending[0].name, "identity__int");

        // ensure cached value is reused and no extra pending instantiations
        monomorphizer
            .instantiate_function("identity", &[Type::Basic(BasicType::Int)])
            .expect("second instantiation should succeed");
        assert!(monomorphizer
            .take_pending_function_instantiations()
            .is_empty());
    }
}
